import torch

from problems.DTLZ import DTLZ


class DTLZ5(DTLZ):
    def __init__(self, var_dim: int, obj_dim: int, max_fun_eval, kwargs: dict):
        super().__init__(var_dim, obj_dim, max_fun_eval, kwargs)

    def eval_value(self, dec):
        g = torch.sum((dec[:, self.obj_dim - 1:] - 0.5) ** 2, dim=1)
        return super().calc_zdt5_6(dec, g)

    def get_optimal_solutions(self, size):
        return super().get_optimal_solutions2(size)
